import numpy as np
from matplotlib import pyplot as plt
import matplotlib.animation as animation

# 基本参数
DATA_SIZE = 300  # 数据集大小
HIST_BINS = np.linspace(0, 30, 21)  # X轴范围

# 二项分布及其参数
DISTRIBUTION_1 = np.random.binomial
N_B = 20  # 实验次数
P_B = 0.6  # 抽中概率

# 泊松分布及其参数
DISTRIBUTION_2 = np.random.poisson
LAM = 15


class MergeDistributions(object):
    def __init__(self) -> None:
        # 生成基本画布
        self.fig, self.ax = plt.subplots()
        # 依照函数生成数列
        self.data_b = DISTRIBUTION_1(N_B, P_B, DATA_SIZE)
        self.data_p = DISTRIBUTION_2(LAM, DATA_SIZE)

    def animate(self, frame_number) -> None:
        """
        绘制一幅画面
        """
        self.ax.clear()
        # 绘制泊松分布
        y_p, bins_p, self.bar_p = self.ax.hist(
            self.data_p[: frame_number + 1],
            bins=HIST_BINS,
            alpha=0.5,
            color='blue',
            label=f'possion({LAM})'
        )
        self.ax.plot(HIST_BINS[:-1], y_p, ls='--')
        # 绘制二项分布
        y_b, bins_b, self.bar_b = self.ax.hist(
            self.data_b[: frame_number + 1],
            bins=HIST_BINS,
            alpha=0.5,
            color='red',
            label=f'binomial({N_B}, {P_B})'
        )
        self.ax.plot(HIST_BINS[:-1], y_b, ls='--')
        # 混合
        mix = y_b * y_p / 100
        self.ax.plot(
            HIST_BINS[:-1],
            mix,
            color='black',
            label='binomial * possion / 100'
        )

        self.ax.grid()
        self.ax.legend()

    def run(self) -> None:
        """
        绘制所有画面形成动图
        """
        ani = animation.FuncAnimation(
            self.fig,
            self.animate,
            DATA_SIZE,
            interval=50,
            repeat=False
        )
        ani.save('binomial_and_poission.gif')


if __name__ == '__main__':
    MergeDistributions().run()
